import torch
from typing import Dict, List
from ..shard.placement_types import DeviceMesh,AbstractTensor,OperationData, OperationDataType
from .node_handler import NodeHandler,FollowHandler
from ..generator import WhereStrategyGenerator,StrategyGenerator
from ..StrategiesVector import StrategiesVector
from geesibling.core.types import Graph, Node
from .registry import operator_registry
__all__ = ['WhereHandler']


@operator_registry.register('aten.where.self')   #arange算子通常用于生成位置编码（Positional Encoding）或者作为序列索引
class WhereHandler(FollowHandler):
    """
    A TensorConstructorHandler which deals with the sharding strategies for tensor constructor operations, such as torch.arange.
    """
    def get_strategy_generator(self) -> List[StrategyGenerator]:
        op_data_mapping = self.get_operation_data_mapping()
        generators = []
        # condition x y
        if self.node.input_shape(1):
            generators.append(WhereStrategyGenerator(op_data_mapping, self.device_mesh,self.pre_strategies_vectors[self.node.inputs[1]]))
        else:
            generators.append(WhereStrategyGenerator(op_data_mapping, self.device_mesh,self.pre_strategies_vectors[self.node.inputs[2]]))
        return generators

    def get_operation_data_mapping(self) -> Dict[str, OperationData]:
        # use transposed shape for strategies
        # the strategies will be transformed back to its original shape in self.post_process
        # TODO 修改所有的data为TensorMeta
        condition_abs_data = AbstractTensor(self.node.input_shape(0),self.node.input_type(0))
        x_abs_data = AbstractTensor(self.node.input_shape(1),self.node.input_type(1))
        y_abs_data = AbstractTensor(self.node.input_shape(1),self.node.input_type(2))
        # if self.node.input_shape(1):
        #     x_abs_data = AbstractTensor(self.node.input_shape(1),self.node.input_type(1))
        #     x_data = OperationData(name=str(self.node.inputs[1]), data=x_abs_data,type=OperationDataType.INPUT)
        # else:
        #     x_abs_data = AbstractTensor(self.node.input_shape(2),self.node.input_type(2))
        #     x_data = OperationData(name=str(self.node.inputs[2]), data=x_abs_data,type=OperationDataType.INPUT)
        output_abs_data = AbstractTensor(self.node.output_shape(0),self.node.output_type(0))
        op_data_mapping = {
            "condition":OperationData(name=str(self.node.inputs[0]), data=condition_abs_data,type=OperationDataType.ARG),
            "x":OperationData(name=str(self.node.inputs[1]), data=x_abs_data,type=OperationDataType.INPUT),
            "y":OperationData(name=str(self.node.inputs[2]), data=y_abs_data,type=OperationDataType.ARG),
            "output":OperationData(name=str(self.node.name), data=output_abs_data,type=OperationDataType.OUTPUT)
        }
        return op_data_mapping
